package com.worley.argutils.bk.command;

import java.util.Arrays;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;
import java.util.TreeSet;
import java.util.stream.Collectors;

import com.worley.argutils.bk.Benchmark;
import com.worley.argutils.bk.RunState;

import redis.clients.jedis.Jedis;
import redis.clients.jedis.JedisCluster;
import redis.clients.jedis.JedisPool;

public class KeysCommand extends Benchmark {

    public static final String NAME = "KEYS";
    
    private JedisCluster jc;

    public KeysCommand(JedisCluster jc, RunState rs) {
        super(rs);
        this.jc = jc;
    }

    @Override
    public void run() {
        Jedis jedis = jc.getClusterNodes().entrySet().iterator().next()
                .getValue().getResource();
        Set<String> masterNodes = Arrays
                .asList(jedis.clusterNodes().split("\n")).stream()
                .filter(l -> l.contains("master"))
                .map(l -> l.split(" ")[1].split("@")[0])
                .collect(Collectors.toSet());

        int loops = rs.getLoops();
        do {
            long s = System.nanoTime();
            for (int i = 0; i < RunState.BATCH_UPDATE_COUNT; i++) {
                keys(masterNodes, loops, i);
            }
            rs.accTotalReq();
            rs.accTotalNanoTime(System.nanoTime() - s);
        } while (--loops > 0);
    }

    private TreeSet<String> keys(Set<String> masterNodes, int loops, int index) {
        String pattern = String.format(KEYS_FORMAT, loops, index);
        TreeSet<String> keys = new TreeSet<>();
        Map<String, JedisPool> clusterNodes = jc.getClusterNodes();

        for (Entry<String, JedisPool> item : clusterNodes.entrySet()) {
            // 是主节点
            if (masterNodes.contains(item.getKey())) {
                try (Jedis jedis = item.getValue().getResource()) {
                    keys.addAll(jedis.keys(pattern));
                } catch (Exception e) {
                    logger.error("Getting keys error: {}", pattern, e);
                }
            }

        }
        return keys;
    }

    @Override
    public String getName() {
        return NAME;
    }

}
